# -*- coding: utf-8 -*-
# @Time    : 2025/1/15 上午11:24
# @Author  : ysj
# @FileName: train_data_split.py
# @Software: PyCharm
# @Blog    ：https://blog.csdn.net/ydscc?type=blog

import os
import random
import numpy as np
import shutil


images_dir = r'C:\Al\Software\AI_Model\Project\yolov5\yolov5\my_data\dataset_20250116_110206_yolo\dataset\images'
labels_dir = r'C:\Al\Software\AI_Model\Project\yolov5\yolov5\my_data\dataset_20250116_110206_yolo\dataset\labels'

images_paths = [os.path.join(images_dir, fname) for fname in os.listdir(images_dir)]
labels_paths = [os.path.join(labels_dir,fname) for fname in os.listdir(labels_dir)]

data_pairs = list(zip(images_paths, labels_paths))

random.shuffle(data_pairs)

spilt_index = int(0.9*len(data_pairs))

train_data = data_pairs[:spilt_index]
val_data = data_pairs[spilt_index:]

os.makedirs('./yolov5_train/train/images', exist_ok=True)
os.makedirs('./yolov5_train/train/labels', exist_ok=True)
os.makedirs('./yolov5_train/val/images', exist_ok=True)
os.makedirs('./yolov5_train/val/labels', exist_ok=True)


for image_path, label_path in train_data:
    shutil.copy(image_path,'./yolov5_train/train/images')
    shutil.copy(label_path, './yolov5_train/train/labels')

for image_path, label_path in val_data:
    shutil.copy(image_path, './yolov5_train/val/images')
    shutil.copy(label_path, './yolov5_train/val/labels')
